import java.util.*;

/**
 * 863. 二叉树中所有距离为 K 的结点
 * https://leetcode-cn.com/problems/all-nodes-distance-k-in-binary-tree/
 */
public class Solutions_863 {
    public static void main(String[] args) {
        String str = "[3 5 1 6 2 0 8 # # 7 4]";
//        int val = 5;
//        int K = 2;  // output: {7, 4, 1}
//        int K = 3;  // output: {8, 0}

        int val = 2;
        int K = 1;  // output: {5, 7, 4}

        TreeNode root = MyTreeNodeUtils.deserialize(str);
        TreeNode target = MyTreeNodeUtils.getNode(root, val);
        List<Integer> result = distanceK(root, target, K);
        System.out.println(result);
    }

    private static Map<TreeNode, TreeNode> parentMap;
    /**
     * 解法二：队列（ms）
     * 将二叉树看作图，而 target 节点作为图的遍历入口
     * 再使用 BFS 来一层一层向外扩散，最终找到全部的与 target 距离为 K 的节点
     */
    public static List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        List<Integer> res = new ArrayList<>();
        parentMap = new HashMap<>();
        // key：节点，value：父节点
        // 方便找到 target 节点的父节点，再遍历其父节点、祖父节点的子节点
        dfsToMap(root, null);
        Queue<TreeNode> queue = new LinkedList<>();
        // 记录下已访问过的节点
        Set<TreeNode> visited = new HashSet<>();
        queue.add(target);
        visited.add(target);

        int dist = 0;
        while (!queue.isEmpty()) {
            if (dist == K) {
                while (!queue.isEmpty()) {
                    // 此时 queue 中的元素，就是与 target 节点，距离为 K 的节点
                    res.add(queue.poll().val);
                }
                return res;
            }
            int size = queue.size();
            for (int i = 0; i < size; i++) {
                TreeNode node = queue.poll();
                // 添加左右子节点
                TreeNode left = node.left, right = node.right;
                if (left != null && !visited.contains(left)) {
                    queue.add(left);
                    visited.add(left);
                }
                if (right != null && !visited.contains(right)) {
                    queue.add(right);
                    visited.add(right);
                }
                // 添加父节点
                TreeNode parent = parentMap.get(node);
                if (parent != null && !visited.contains(parent)) {
                    queue.add(parent);
                    visited.add(parent);
                }
            }
            // 距离长度加 1
            dist ++;
        }
        return new ArrayList<>();
    }

    public static void dfsToMap(TreeNode node, TreeNode parent) {
        if (node == null) {
            return;
        }
        if (parent != null) {
            parentMap.put(node, parent);
        }
        dfsToMap(node.left, node);
        dfsToMap(node.right, node);
    }



    private static int targetVal;
    private static List<Integer> res;
    private static int k;

    /**
     * 解法一：计算距离（18ms）
     * 注意：target 节点的每个父节点的左/右子树下，都有可能存在距离为 K 的节点
     */
    public static List<Integer> distanceK2(TreeNode root, TreeNode target, int K) {
        res = new ArrayList<>();
//        targetVal = target.val;
        targetVal = target.val;
        k = K;
        dfs(root);
        return res;
    }

    public static int dfs(TreeNode node) {
        if (node == null) {
            return -1;
        }
        if (node.val == targetVal) {
            // 找到 target 节点，那么在 target 的子树中，找到距离为 k 的子节点
            getDistNode(node, k);
            // 关键：返回 1，即 node 节点与父节点之间的距离为 1
            return 1;
        } else {
            // target 若在左子树中，那么距离 node 的节点距离为 leftRes，为 -1 说明不在左子树中
            int leftRes = dfs(node.left);
            if (leftRes > 0) {
                int dist = k - leftRes;
                if (dist == 0) {
                    // 不需要遍历 node 的子树
                    res.add(node.val);
                } else if (dist > 0) {
                    // 关键：防止重复添加，遍历 node 的右子树，在右子树中找到距离为 dist - 1 的节点
                    getDistNode(node.right, dist - 1);
                }
                return leftRes + 1;
            }
            // 若 target 已经存在于左子树中，那么不再递归查找右子树
            // target 若在右子树中，那么距离 node 的节点距离为 rightRes，为 -1 说明不在右子树中
            int rightRes = dfs(node.right);
            if (rightRes > 0) {
                int dist = k - rightRes;
                if (dist == 0) {
                    // 不需要遍历 node 的子树
                    res.add(node.val);
                } else if (dist > 0) {
                    // 关键：防止重复添加，遍历 node 的左子树，在左子树中找到距离为 dist - 1 的节点
                    getDistNode(node.left, dist - 1);
                }
                return rightRes + 1;
            }
        }
        return -1;
    }

    public static void getDistNode(TreeNode node, int len) {
        if (node == null) {
            return;
        }
        if (len == 0) {
            // 剪枝：找到距离相等的节点，不需要再继续递归
            res.add(node.val);
            return;
        }
        // 遍历左右节点，路径长度减 1
        getDistNode(node.left, len - 1);
        getDistNode(node.right, len - 1);
    }
}
